package semantic_compute.entity_linking.yago_linker;

import algorithm.kg.mpi_aida.AIDAClient;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ArrayNode;
import common.helper.json.JsonToJavaObject;
import semantic_compute.entity_linking.schema.EntityLinkingResult;

import java.util.*;

/**
 * Created by julianzliu on 4/16/2017.
 */
public class YagoTextsEntityLinker {

    private AIDAClient aidaClient;

    public YagoTextsEntityLinker(){
        this.aidaClient = new AIDAClient();
    }


    public EntityLinkingResult getYagoEntityLinkingResult(String text){
        JsonNode result = getYagoEntityLinkingJson( text );
        String petty_reuslt = JsonToJavaObject.jsonNodeToString(result);
        System.out.println( petty_reuslt );
        EntityLinkingResult linkingResult = parseEntityLinkingResult(result);
        //System.out.println("[getYagoEntityLinkingResult] parseEntityLinkingResult Finished!");

        try {
            /*
            TODO:
             */
            //linkingResult.extractLinkedEntities();
        }catch (Exception e){
            System.out.println(e.getMessage());
        }


        return linkingResult;
    }


    public JsonNode getYagoEntityLinkingJson(String text){
        String json_result = this.aidaClient.entity_linking( text );
        System.out.println("[Raw Entity Linking Result]\n" + json_result + "\n\n");
        JsonNode reuslt = JsonToJavaObject.loadJsonNode(json_result);
        return reuslt;
    }


    public EntityLinkingResult parseEntityLinkingResult(JsonNode result){
        //System.out.println("[parseEntityLinkingResult] begin parse");
        EntityLinkingResult linkingResult = new EntityLinkingResult();
        linkingResult.originalText = result.get("originalText").asText();
        linkingResult.annotatedText = result.get("annotatedText").asText();


        // entityMetadata
        //System.out.println("[parseEntityLinkingResult] begin parse Metadata");
        Map<String,EntityLinkingResult.EntityMetadata> entityMetadataMap = new HashMap<>();
        JsonNode entityMetadatas = result.get("entityMetadata");
        //Iterable<StirngentityMetadatas.fieldNames();
        Iterator<String> entityNames = entityMetadatas.fieldNames();
        while( entityNames.hasNext() ){
            //entityNames.
            String entity_name = entityNames.next();
            if(entity_name.equalsIgnoreCase("AIDA:--OOKBE--"))
                continue;
            //System.out.println("  [parseEntityLinkingResult] get Name " + entity_name);
            JsonNode entity = entityMetadatas.get(entity_name);
            EntityLinkingResult.EntityMetadata metadata = new EntityLinkingResult.EntityMetadata();
            metadata.knowledgebase = entity.get("knowledgebase").asText();
            metadata.entityId = "<" + entity.get("entityId").asText() + ">";
            metadata.YAGO_ID = metadata.entityId;
            metadata.url = entity.get("url").asText();
            metadata.importance =entity.get("importance").asDouble();
            entityMetadataMap.put(entity_name, metadata);
        }
        linkingResult.entityMetadata = entityMetadataMap;

        // mentions
        System.out.println("[parseEntityLinkingResult] begin parse mentions");
        List<EntityLinkingResult.LinkedEntity> entities = new ArrayList<>();
        ArrayNode mentions = (ArrayNode)result.get("mentions");
        for(JsonNode mention : mentions){
            EntityLinkingResult.LinkedEntity entity = new EntityLinkingResult.LinkedEntity();
            entity.name = mention.get("name").asText();
            entity.offset = mention.get("offset").asInt();
            entity.length = mention.get("length").asInt();
            if(mention.hasNonNull("bestEntity")){
                entity.entity = mention.get("bestEntity").get("kbIdentifier").asText();
                entity.score = mention.get("bestEntity").get("disambiguationScore").asDouble();
            }
            else {
                String[] names = entity.name.split(" ");
                String entity_str = "<";
                for(int i = 0; i < names.length-1; i++) entity_str += names[i] + "_";
                entity_str += names[names.length-1] + ">";
                entity.entity = entity_str;
                entity.score = 0.1;
            }

            entities.add(entity);
        }
        linkingResult.mentions = entities;

        System.out.println("[parseEntityLinkingResult] end parse");

        return linkingResult;
    }




    /******************************************************
     * MAIN
     */
    public static void main(String[] args){
        String text;
        //text = "Obama was born in Hawaii. He is t6he president of USA.";
        //text = "where Obama president was born?";
        text = "population of Seattle. The Weather of Seattle is rainy.";

        text = "Craig Malcolm Robinson is elder brother of Michelle Obama.";
        YagoTextsEntityLinker entityLinker = new YagoTextsEntityLinker();




        JsonNode json_result = entityLinker.getYagoEntityLinkingJson( text );
        String petty_reuslt = JsonToJavaObject.jsonNodeToString(json_result);
        System.out.println( petty_reuslt );


        EntityLinkingResult result = entityLinker.getYagoEntityLinkingResult( text );
        result.alignWithCoreNlp();

        System.out.println("\n\n\n\n--------------------------------");
        System.out.println(result.toString());


    }


}
